import torch as pt
import torchvision as ptv
import torch.nn as nn
import numpy as np
from prepare_data import GL_FLATTEN_DEPTH, GL_NUMGRID, GL_NUMBBOX, GL_N_CLS
from util_again import calculate_iou
from python_ai.common.xcommon import *


class MyNet(nn.Module):

    def __init__(self):
        super().__init__()
        resnet = ptv.models.resnet34(pretrained=True)

        resnet_output_ch = resnet.fc.in_features

        self.resnet = nn.Sequential(*(list(resnet.children())[:-2]))
        self.convs = nn.Sequential(
            nn.Conv2d(resnet_output_ch, 1024, 3, padding=1),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(inplace=True),

            nn.Conv2d(1024, 1024, 3, padding=1, stride=2),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(inplace=True),

            nn.Conv2d(1024, 1024, 3, padding=1),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(inplace=True),

            nn.Conv2d(1024, 1024, 3, padding=1),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(inplace=True),
        )
        self.conns = nn.Sequential(
            nn.Conv2d(1024, GL_FLATTEN_DEPTH, 3, padding=1),
            nn.BatchNorm2d(GL_FLATTEN_DEPTH),
            nn.Sigmoid(),
        )

    def forward(self, inputs):
        x = self.resnet(inputs)
        print('after resnet x.size:', x.size())
        x = self.convs(x)
        x = self.conns(x)
        self.pred = x
        return x

    def calculate_loss(self, labels):
        lambda_coord, lambda_nonobj = 5., 0.5
        loss_location, loss_proba, loss_cls = 0., 0., 0.
        n_samples = len(self.pred)
        for sample_id in range(n_samples):
            for m in range(GL_NUMGRID):
                for n in range(GL_NUMGRID):
                    if np.isclose(0., labels[sample_id, 4, m, n]):
                        loss_proba += lambda_nonobj * pt.sum(self.pred[sample_id, [4, 9], m, n] ** 2)
                    else:
                        loss_cls += pt.sum((self.pred[sample_id, 10:, m, n] - labels[sample_id, 10:, m, n]) ** 2)

                        # iou
                        box1_xy = (
                            (self.pred[sample_id, 0, m, n] + n) / GL_NUMGRID + self.pred[sample_id, 2, m, n] / 2,
                            (self.pred[sample_id, 1, m, n] + m) / GL_NUMGRID + self.pred[sample_id, 3, m, n] / 2,
                            (self.pred[sample_id, 0, m, n] + n) / GL_NUMGRID - self.pred[sample_id, 2, m, n] / 2,
                            (self.pred[sample_id, 1, m, n] + m) / GL_NUMGRID - self.pred[sample_id, 3, m, n] / 2,
                        )
                        box2_xy = (
                            (self.pred[sample_id, 5, m, n] + n) / GL_NUMGRID + self.pred[sample_id, 7, m, n] / 2,
                            (self.pred[sample_id, 6, m, n] + m) / GL_NUMGRID + self.pred[sample_id, 8, m, n] / 2,
                            (self.pred[sample_id, 5, m, n] + n) / GL_NUMGRID - self.pred[sample_id, 7, m, n] / 2,
                            (self.pred[sample_id, 6, m, n] + m) / GL_NUMGRID - self.pred[sample_id, 8, m, n] / 2,
                        )
                        gt_xy = (
                            (labels[sample_id, 0, m, n] + n) / GL_NUMGRID + labels[sample_id, 2, m, n] / 2,
                            (labels[sample_id, 1, m, n] + m) / GL_NUMGRID + labels[sample_id, 3, m, n] / 2,
                            (labels[sample_id, 0, m, n] + n) / GL_NUMGRID - labels[sample_id, 2, m, n] / 2,
                            (labels[sample_id, 1, m, n] + m) / GL_NUMGRID - labels[sample_id, 3, m, n] / 2,
                        )
                        iou1 = calculate_iou(box1_xy, gt_xy)
                        iou2 = calculate_iou(box2_xy, gt_xy)

                        if iou1 > iou2:
                            loss_location += lambda_coord * pt.sum(
                                (self.pred[sample_id, 0:2, m, n] - labels[sample_id, 0:2, m, n]) ** 2
                                + (self.pred[sample_id, 2:4, m, n].sqrt() - labels[sample_id, 2:4, m, n].sqrt()) ** 2
                            )
                            loss_proba += pt.sum((self.pred[sample_id, 4, m, n] - iou1) ** 2)
                            loss_proba += lambda_nonobj * pt.sum((self.pred[sample_id, 9, m, n] - iou2) ** 2)
                        else:
                            loss_location += lambda_coord * pt.sum(
                                    (self.pred[sample_id, 5:7, m, n] - labels[sample_id, 5:7, m, n]) ** 2
                                    + (self.pred[sample_id, 7:9, m, n].sqrt() - labels[sample_id, 7:9, m, n].sqrt()) ** 2
                            )
                            loss_proba += pt.sum((self.pred[sample_id, 9, m, n] - iou2) ** 2)
                            loss_proba += lambda_nonobj * pt.sum((self.pred[sample_id, 4, m, n] - iou1) ** 2)
        loss = loss_location + loss_proba + loss_cls
        loss /= n_samples
        return loss


if __name__ == '__main__':
    # 自定义输入张量，验证网络可以正常跑通，并计算loss，调试用
    x = pt.zeros(5,3,448,448)
    net = MyNet()
    sep('net')
    print(net)
    a = net(x)
    print(a.shape)

    # test with objs and non-objs
    label_file_path = r'../../../../../large_data\CV3\_many_files\yolov1_person_car/voc2012_forYolov1/train.csv'
    labels = np.loadtxt(label_file_path)[:5].reshape(-1, 7, 7, GL_FLATTEN_DEPTH)
    print_numpy_ndarray_info(labels, 'labels')
    labels = np.transpose(labels, [0, 3, 1, 2])
    print_numpy_ndarray_info(labels, 'labels')
    labels = pt.Tensor(labels)

    loss = net.calculate_loss(labels)
    print(loss)
